from config import CLASS_NUM
from torchvision.models import vgg16
import torch
import torch.nn as nn

def vgg16_101(**kwargs):
    net = vgg16(**kwargs)
    for params in net.parameters():
        params.requires_grad = False
    net.classifier[6] = nn.Sequential(
        nn.Linear(4096, 256),
        nn.ReLU(inplace=True),
        nn.Dropout(0.4),
        nn.Linear(256, CLASS_NUM),
    )
    
    return net

if __name__ == "__main__":
    net = vgg16_101(pretrained=True)
    dummy = torch.rand((4, 3, 256, 256))
    result = net(dummy)
    print(result.size())